package com.sjr.easy.chat.utils;

import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;

import java.util.UUID;

public class TraceUtils {

    public static final String HTTP_TRACE_ID_KEY = "__HTTP_TRACE_ID__";
    public static final String RABBIT_MQ_TRACE_ID_KEY = "__RABBIT_MQ_TRACE_ID__";

    public static final String TRACE_ID_KEY = "traceId";
    public static final String SPAN_ID_KEY = "spanId";

    public static <T> T runWithTracer(TracerCallable<T> callable, TraceInfo traceInfo) throws Throwable {
        if (StringUtils.isEmpty(traceInfo.traceId) || StringUtils.isEmpty(traceInfo.spanId)) {
            return callable.call();
        } else {
            try {
                setTraceId(traceInfo.traceId);
                setSpanId(traceInfo.spanId);
                return callable.call();
            } finally {
                cleanTraceInfo();
            }
        }
    }

    public static void setTraceId(String traceId) {
        MDC.put(TRACE_ID_KEY, traceId);
    }


    public static String currentTraceId() {
        return StringUtils.defaultString(MDC.get(TRACE_ID_KEY), null);
    }

    public static void setSpanId(String spanId) {
        MDC.put(SPAN_ID_KEY, spanId);
    }

    public static String currentSpanId() {
        return StringUtils.defaultString(MDC.get(SPAN_ID_KEY), null);
    }

    public static void cleanTraceInfo() {
        MDC.remove(TRACE_ID_KEY);
        MDC.remove(SPAN_ID_KEY);
    }

    public static String generateId() {
        return UUID.randomUUID().toString().replaceAll("-", "");
    }

    public static class TraceInfo {
        public String traceId;
        public String spanId;

        public TraceInfo() {
        }

        public TraceInfo(String traceId, String spanId) {
            this.traceId = traceId;
            this.spanId = spanId;
        }
    }

    public interface TracerCallable<V> {
        V call() throws Throwable;
    }


}
